import h5py
import torch
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
import numpy as np
import sys
sys.path.append('..')
from DataTransers.TraceTranser import *

class Datasetloader():
    def __init__(self, config):
        self.datasets_config = config["datasets"]
        self.label_type = self.datasets_config["label_type"]

    def get_train_loader(self):
        # 读取配置
        self.cfg_trainset = self.datasets_config["trainset"]
        bs = self.cfg_trainset["batch_size"]
        shuffle_ = self.cfg_trainset["shuffle"]
        left = self.cfg_trainset["left"]
        right = self.cfg_trainset["right"]
        in_file = h5py.File(self.cfg_trainset["path"], 'r')
        
        # 读取数据
        traces_profiling = np.array(in_file['Profiling_traces/traces'][left:right], dtype=np.float32)
        labels_profiling = np.array(in_file['Profiling_traces/labels'][left:right])
        # 存储不同变换曲线数据集
        traces_list = []
        trans_traces = self.cfg_trainset["trans_traces"]
        for trace_cfg in trans_traces:
            trans_queue = trace_cfg["trans_queue"]
            if (trans_queue == None):
                trans_queue = []
            traces_profiling_n = traces_profiling
            for tran_method in trans_queue:
                if (tran_method == "loopmove"):
                    traces_profiling_n = LoopMove(trace_cfg["move"])(traces_profiling_n)
                elif (tran_method == "randomloopmove"):
                    traces_profiling_n = RandomLoopMove(trace_cfg["randm_min"], trace_cfg["randm_max"])(traces_profiling_n)
                elif (tran_method == "integrate"):
                    traces_profiling_n = Integrate(trace_cfg["in_step"], trace_cfg["in_mode"])(traces_profiling_n)
                elif (tran_method == "gaunoise"):
                    traces_profiling_n = AddGaussianNoise(trace_cfg["gau_mean"], trace_cfg["gau_std"])(traces_profiling_n)
                elif (tran_method == "normal"):
                    traces_profiling_n = Normalize()(traces_profiling_n)
                elif (tran_method == "horflip"):
                    traces_profiling_n = HorizontalFlip()(traces_profiling_n)
                elif (tran_method == "pca"):
                    traces_profiling_n = PCATransform(trace_cfg["pca_dim"])(traces_profiling_n)
            traces_list.append(torch.from_numpy(traces_profiling_n))
            # 曲线处理完毕

        # 输出labels_profiling的形状
        labels = torch.from_numpy(labels_profiling)
        labels = self.convert_labels(labels, self.label_type)

        # 生成数据集
        ascad_dataset = TensorDataset(*traces_list, labels)
        loader = DataLoader(dataset=ascad_dataset, batch_size=bs, shuffle=shuffle_)
        return loader

    def get_val_loader(self):
        # 读取配置
        self.cfg_valset = self.datasets_config["valset"]
        bs = self.cfg_valset["batch_size"]
        left = self.cfg_valset["left"]
        right = self.cfg_valset["right"]
        # 读取数据
        if (self.cfg_valset["path"] == ""):
            in_file = h5py.File(self.cfg_trainset["path"], 'r')
            traces_attack = np.array(in_file['Profiling_traces/traces'][left:right], dtype=np.float32)
            labels_attack = np.array(in_file['Profiling_traces/labels'][left:right])
            metadatas = in_file['Profiling_traces/metadata'][left:right]
            plain_text = []
            for md in metadatas:
                plain_text.append(md[0][2])
            plain_text_ = torch.from_numpy(np.array(plain_text))
        else:
            in_file = h5py.File(self.cfg_valset["path"], 'r')
            traces_attack = np.array(in_file['Attack_traces/traces'][left:right], dtype=np.float32)
            labels_attack = np.array(in_file['Attack_traces/labels'][left:right])
            metadatas = in_file['Attack_traces/metadata'][left:right]
            plain_text = []
            for md in metadatas:
                plain_text.append(md[0][2])
            plain_text_ = torch.from_numpy(np.array(plain_text))
        
        # 存储不同变换曲线数据集
        trans_queue = self.cfg_valset["trans_queue"]
        if (trans_queue == None):
            trans_queue = []
        for tran_method in trans_queue:
            if (tran_method == "loopmove"):
                traces_attack = LoopMove(self.cfg_valset["move"])(traces_attack)
            elif (tran_method == "randomloopmove"):
                traces_attack = RandomLoopMove(self.cfg_valset["randm_min"], self.cfg_valset["randm_max"])(traces_attack)
            elif (tran_method == "integrate"):
                traces_attack = Integrate(self.cfg_valset["in_step"], self.cfg_valset["in_mode"])(traces_attack)
            elif (tran_method == "gaunoise"):
                traces_attack = AddGaussianNoise(self.cfg_valset["gau_mean"], self.cfg_valset["gau_std"])(traces_attack)
            elif (tran_method == "normal"):
                traces_attack = Normalize()(traces_attack)
            elif (tran_method == "horflip"):
                traces_attack = HorizontalFlip()(traces_attack)
            elif (tran_method == "pca"):
                traces_attack = PCATransform(self.cfg_valset["pca_dim"])(traces_attack)
        traces_attack = torch.from_numpy(traces_attack)
        # 曲线处理完毕
        # 标签
        labels = torch.from_numpy(labels_attack)
        labels = self.convert_labels(labels, self.label_type)

        # 生成数据集
        ascad_dataset = TensorDataset(traces_attack, plain_text_, labels)
        loader = DataLoader(dataset=ascad_dataset, batch_size=bs, shuffle=False)
        return loader

    # 计算一个值的汉明重量
    def hamming_weight(self, num):
        count = 0
        while num:
            count += 1
            num &= num - 1
        return count

    # 将labels值转为其他分类的labels
    def convert_labels(self, labels, label_type):
        if (label_type == 256): 
            return labels
        for i, label in enumerate(labels):
            if (label_type == 2): # MSB
                labels[i] = label >> 7
            elif (label_type == 9):
                labels[i] = self.hamming_weight(label.item())
        return labels
